import pandas as pd
import os
from datetime import datetime
import sys
import argparse

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from AI_train.predict_stocks import predict_stocks_for_date

def get_all_2025_dates():
    """
    从CSV文件中提取所有2025年的日期
    """
    # 获取tdx_data目录下的所有CSV文件
    data_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tdx_data')
    csv_files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
    
    # 存储所有2025年的日期
    dates_2025 = set()
    
    # 遍历所有CSV文件，查找2025年的日期
    for csv_file in csv_files:
        file_path = os.path.join(data_dir, csv_file)
        try:
            df = pd.read_csv(file_path)
            # 获取第一列的列名（日期列）
            date_column = df.columns[0]
            # 筛选2025年的数据
            df_2025 = df[df[date_column].str.startswith('2025')]
            for date in df_2025[date_column]:
                dates_2025.add(date)
        except Exception as e:
            print(f"读取文件 {csv_file} 时出错: {e}")
    
    # 将日期按时间顺序排序
    sorted_dates = sorted(list(dates_2025))
    return sorted_dates

def get_dates_in_range(start_date=None, end_date=None):
    """
    从CSV文件中提取指定日期范围内的所有日期
    
    Parameters:
    start_date: str, 开始日期，格式为'YYYY-MM-DD'
    end_date: str, 结束日期，格式为'YYYY-MM-DD'
    """
    # 获取tdx_data目录下的所有CSV文件
    data_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tdx_data')
    csv_files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
    
    # 存储所有日期
    all_dates = set()
    
    # 遍历所有CSV文件，查找日期
    for csv_file in csv_files:
        file_path = os.path.join(data_dir, csv_file)
        try:
            df = pd.read_csv(file_path)
            # 获取第一列的列名（日期列）
            date_column = df.columns[0]
            
            # 获取所有日期
            for date in df[date_column]:
                all_dates.add(date)
        except Exception as e:
            print(f"读取文件 {csv_file} 时出错: {e}")
    
    # 转换为列表并排序
    sorted_dates = sorted(list(all_dates))
    
    # 如果提供了日期范围，则进行筛选
    if start_date or end_date:
        # 转换为datetime对象进行比较
        start_dt = datetime.strptime(start_date, '%Y-%m-%d') if start_date else None
        end_dt = datetime.strptime(end_date, '%Y-%m-%d') if end_date else None
        
        # 筛选日期范围内的日期
        filtered_dates = []
        for date_str in sorted_dates:
            date_dt = datetime.strptime(date_str, '%Y-%m-%d')
            if start_dt and date_dt < start_dt:
                continue
            if end_dt and date_dt > end_dt:
                continue
            filtered_dates.append(date_str)
        
        sorted_dates = filtered_dates
    
    return sorted_dates

def predict_all_2025_dates():
    """
    预测所有2025年的日期
    """
    # 获取所有2025年的日期
    dates_2025 = get_all_2025_dates()
    
    print(f"找到 {len(dates_2025)} 个2025年的日期")
    
    # 对每个日期进行预测
    for date in dates_2025:
        print(f"\n正在预测日期: {date}")
        try:
            predict_stocks_for_date(date)
        except Exception as e:
            print(f"预测日期 {date} 时出错: {e}")

def predict_dates_in_range(start_date=None, end_date=None):
    """
    预测指定日期范围内的所有日期
    
    Parameters:
    start_date: str, 开始日期，格式为'YYYY-MM-DD'
    end_date: str, 结束日期，格式为'YYYY-MM-DD'
    """
    # 获取指定范围内的日期
    dates = get_dates_in_range(start_date, end_date)
    
    print(f"找到 {len(dates)} 个日期")
    
    # 对每个日期进行预测
    for date in dates:
        print(f"\n正在预测日期: {date}")
        try:
            predict_stocks_for_date(date)
        except Exception as e:
            print(f"预测日期 {date} 时出错: {e}")

def main():
    """
    主函数
    """
    parser = argparse.ArgumentParser(description='预测指定日期范围内的股票')
    parser.add_argument('--start-date', '--start_date', type=str, help='开始日期 (YYYY-MM-DD)')
    parser.add_argument('--end-date', '--end_date', type=str, help='结束日期 (YYYY-MM-DD)')
    
    args = parser.parse_args()
    
    if args.start_date and args.end_date:
        print(f"开始预测 {args.start_date} 到 {args.end_date} 之间的日期...")
        predict_dates_in_range(args.start_date, args.end_date)
    elif args.start_date:
        print(f"开始预测 {args.start_date} 之后的所有日期...")
        predict_dates_in_range(start_date=args.start_date)
    elif args.end_date:
        print(f"开始预测 {args.end_date} 之前的所有日期...")
        predict_dates_in_range(end_date=args.end_date)
    else:
        print("开始预测所有2025年的日期...")
        predict_all_2025_dates()
    
    print("所有日期预测完成。")

if __name__ == "__main__":
    main()